package com.linkhub.oms.component.trace.util;

import com.linkhub.oms.component.trace.core.QtTraceContext;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import org.slf4j.MDC;

import static com.linkhub.oms.common.constant.TraceConstant.REQUEST_LOG_ID;
import static com.linkhub.oms.common.constant.TraceConstant.SPAN_TRACE_ID;
import static com.linkhub.oms.common.constant.TraceConstant.TRACE_ID;

/**
 * @Description: MDCUtil
 * @Author: echo
 * @Date: 2022/8/31 20:35
 * @Version: 1.0
 */
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class MDCUtil {


    public static void setRequestId(String requestId) {
        MDC.put(REQUEST_LOG_ID, requestId);
    }

    public static String getRequestId() {
        return MDC.get(REQUEST_LOG_ID);
    }

    public static void setSpanTraceId(String traceId) {
        MDC.put(SPAN_TRACE_ID, traceId);
    }

    public static void setTraceId(String traceId) {
        MDC.put(TRACE_ID, traceId);

    }

    public static String getTraceId() {
        return MDC.get(TRACE_ID);
    }

    public static void removeSpanTraceId() {
        MDC.remove(SPAN_TRACE_ID);
    }

    public static void removeAllTraceId() {
        MDC.remove(TRACE_ID);
        MDC.remove(SPAN_TRACE_ID);
    }

    /**
     * 往 {@link MDC} 设置 {@code traceId}, 默认会递增 {@code spanId}
     *
     * @return traceId
     */
    public static String getAndSetTraceId() {
        return getAndSetTraceId(true);
    }


    /**
     * 往 {@link MDC} 设置 {@code traceId}
     *
     * @param incrementSpanId 是否递增spanId
     * @return traceId
     */
    public static String getAndSetTraceId(boolean incrementSpanId) {
        return QtTraceContext.getAndConsumerTraceId(MDCUtil::setTraceId, incrementSpanId);
    }

    /**
     * 往 {@link MDC} 设置 {@code traceId}
     * 一般用于即将跨线程服务调用
     *
     * @return traceId
     */
    public static String getAndSetSpanTraceId() {
        return getAndSetSpanTraceId(true);
    }

    /**
     * 往 {@link MDC} 设置 {@code traceId}
     * 一般用于即将跨线程服务调用
     *
     * @param incrementSpanId 是否递增spanId
     * @return traceId
     */
    public static String getAndSetSpanTraceId(boolean incrementSpanId) {
        return QtTraceContext.getAndConsumerTraceId(MDCUtil::setSpanTraceId, incrementSpanId);
    }

    public static void setTraceAndRequestId() {
        MDCUtil.setTraceId(QtTraceContext.getTraceId());
        // 兼容旧逻辑
        MDCUtil.setRequestId(QtTraceContext.getRequestId());
    }

    public static void createAndSetTraceIdAndRequestId() {
        MDCUtil.setTraceId(QtTraceContext.createTraceId());
        // 兼容旧逻辑
        MDCUtil.setRequestId(QtTraceContext.getRequestId());
    }

    /**
     * 回放 {@code spanTraceId}
     *
     * @return spanTraceId
     */
    public static String replaySpanTraceId() {
        return QtTraceContext.getAndConsumerTraceId(MDCUtil::setSpanTraceId, false);
    }


}
